// DnsProxy.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"


#include    <winsock2.h>
#include    <Ws2tcpip.h>
#include    <stdio.h>
#include    <map>
#include    <string>

using namespace std;

#define     MAXSOCKET       20
#define     MAXDELAY       2000
#define     IP6HOSTFILE     "ip6hosts.txt"
#define     DNSSUFFIX     ".suryani.cn"
#define     DNSSERVER      "172.16.0.3"

typedef struct _UDPSOCK_INFO
{
	bool					InUsed;
	SOCKET                  UDPSocket;
	SOCKADDR_STORAGE        UDPFrom;
	SOCKADDR_IN             UDPTo;
	char                    SendBuffer[1024];
	int                     SendBufferLen;
	int                     gettimes;
	DWORD                   start, end;
}UDPSOCK_INFO, *PUDPSOCK_INFO;

//DNS header structure
struct DNS_HEADER
{
	unsigned	short id;		    // identification number

	unsigned	char rd     :1;		// recursion desired
	unsigned	char tc     :1;		// truncated message
	unsigned	char aa     :1;		// authoritive answer
	unsigned	char opcode :4;	    // purpose of message
	unsigned	char qr     :1;		// query/response flag

	unsigned	char rcode  :4;	    // response code
	unsigned	char cd     :1;	    // checking disabled
	unsigned	char ad     :1;	    // authenticated data
	unsigned	char z      :1;		// its z! reserved
	unsigned	char ra     :1;		// recursion available

	unsigned    char q_count_h;	    // number of question entries
	unsigned    char q_count_l;	    // number of question entries
	unsigned	char ans_count_h;	// number of answer entries
	unsigned	char ans_count_l;	// number of answer entries
	unsigned	char auth_count_h;	// number of authority entries
	unsigned	char auth_count_l;	// number of authority entries
	unsigned	char add_count_h;	// number of resource entries
	unsigned	char add_count_l;	// number of resource entries
};

#pragma pack(push, 1)
struct RES_RECORD
{
	char                name[2];
	unsigned char       type_h;
	unsigned char       type_l;
	unsigned char       _class_h;
	unsigned char       _class_l;
	unsigned char       ttl_hh;
	unsigned char       ttl_h;
	unsigned char       ttl_l;
	unsigned char       ttl_ll;
	unsigned char       data_len_h;
	unsigned char       data_len_l;
	in6_addr            i6addr;
};
#pragma pack(pop)

#pragma pack(push, 1)
struct CNAME_RECORD
{
	char                name[2];
	unsigned char       type_h;
	unsigned char       type_l;
	unsigned char       _class_h;
	unsigned char       _class_l;
	unsigned char       ttl_hh;
	unsigned char       ttl_h;
	unsigned char       ttl_l;
	unsigned char       ttl_ll;
	unsigned char       data_len_h;
	unsigned char       data_len_l;
	char            cname[100];
};
#pragma pack(pop)

//methods
void CreateIP6Table ();
int CheckIP6 ( char*, in6_addr* );
void ReadRequest();
void ReadResponse();
void SendBackToClient();
void CloseSocket(UDPSOCK_INFO* info);

//globals
std::map <string, string>  ip6Table;
UDPSOCK_INFO            UDPSock[MAXSOCKET];
SOCKET                  localDnsSocket;
struct fd_set           fds_read, fds_write;

int main(int argc, char **argv)
{
	WSADATA                 wsaData;
	SOCKADDR_IN             localDnsAddr;
	SOCKADDR_IN             remoteDnsAddr;
	struct timeval          timeout;
	int                     selectret, i;

	//local ipv6 table
	CreateIP6Table();

	// Initialize Winsock version 2.2
	if( WSAStartup(MAKEWORD(2,2), &wsaData) != 0)
	{
		printf("Server: WSAStartup failed with error %ld\n", WSAGetLastError());
		return -1;
	}

	//init local socket
	localDnsSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (localDnsSocket == INVALID_SOCKET)
	{
		printf("Server: Error at udp socket(): %ld\n", WSAGetLastError());
		WSACleanup();
		return -1;
	}

	//bind local socket
	localDnsAddr.sin_family = AF_INET;
	localDnsAddr.sin_port = htons(53);
	localDnsAddr.sin_addr.s_addr = htonl(INADDR_ANY);

	if (bind(localDnsSocket, (SOCKADDR *)&localDnsAddr, sizeof(localDnsAddr)) == SOCKET_ERROR)
	{
		printf("Server: udp bind() failed! Error: %ld.\n", WSAGetLastError());
		closesocket(localDnsSocket);
		WSACleanup();
		return -1;
	}

	printf("DNS Proxy Starting...\n");

	//init remote socket info
	remoteDnsAddr.sin_family = AF_INET;
	remoteDnsAddr.sin_addr.s_addr =inet_addr( DNSSERVER );
	remoteDnsAddr.sin_port = htons( 53 );

	for ( i=0; i<MAXSOCKET; i++)
	{
		UDPSock[i].UDPSocket = INVALID_SOCKET;
		UDPSock[i].UDPTo = remoteDnsAddr;
		UDPSock[i].gettimes=0;
		UDPSock[i].start=0;
		UDPSock[i].end=0;
		UDPSock[i].InUsed=false;
	}

	//init time out
	timeout.tv_sec = 0;
	timeout.tv_usec = 100;

	while (true)
	{
		FD_ZERO(&fds_read);
		FD_ZERO(&fds_write);
		FD_SET(localDnsSocket, &fds_read);

		for (i=0; i<MAXSOCKET; i++)
		{
			if (UDPSock[i].UDPSocket != INVALID_SOCKET)
				FD_SET(UDPSock[i].UDPSocket, &fds_read);
		}

		selectret = select(0, &fds_read, &fds_write, 0, &timeout);

		switch (selectret)
		{
		case 0:
			// Timed out, do whatever you want to handle this situation
			//printf("Server: Timeout while waiting !...\n");
			break;

		case SOCKET_ERROR:
			printf("Server: Some error encountered with code number: %ld\n", WSAGetLastError());
			break;

		default:
			ReadRequest();

			ReadResponse();
			break;        
		}

		//send back to client
		SendBackToClient();
	}        

	printf("Server: Finished receiving. Closing the socket...\n");

	if (closesocket(localDnsSocket) != 0)
		printf("Server: udp closesocket() failed! Error code: %ld\n", WSAGetLastError());
	else
		printf("Server: udp closesocket() is OK...\n");

	printf("Server: Cleaning up...\n");

	if(WSACleanup() != 0)
		printf("Server: WSACleanup() failed! Error code: %ld\n", WSAGetLastError());
	else
		printf("Server: WSACleanup() is OK\n");

	return 0;
}

void CloseSocket(UDPSOCK_INFO* info)
{
	if (info->UDPSocket != INVALID_SOCKET)
	{
		closesocket(info->UDPSocket);
	}
	info->UDPSocket=INVALID_SOCKET;
	info->InUsed=false;
	info->gettimes = 0;
	info->start = 0;
	info->end = 0;
}

void SendBackToClient()
{
	for (int i=0; i<MAXSOCKET; i++)
	{
		if(!UDPSock[i].InUsed) continue;

		DWORD tmnow = GetTickCount();
		DWORD delay = tmnow -  UDPSock[i].start;

		if (UDPSock[i].gettimes > 0)
		{
			int ret = sendto(localDnsSocket, UDPSock[i].SendBuffer, UDPSock[i].SendBufferLen, 0, (SOCKADDR *)&UDPSock[i].UDPFrom, sizeof(SOCKADDR_STORAGE));

			//print log
			if (UDPSock[i].gettimes == 100)
				printf("used %dms, !! local force ipv6 !!\t", delay);
			else    
				printf("used %dms, got %d UDP(s) in %dms\t", delay, UDPSock[i].gettimes, UDPSock[i].end - UDPSock[i].start );
			int j = 13;
			while (UDPSock[i].SendBuffer[j] != 0 )
			{
				if (UDPSock[i].SendBuffer[j]<0x20) printf("."); else printf("%c", UDPSock[i].SendBuffer[j]);
				j++;
			}    
			printf("\n");

			CloseSocket(&UDPSock[i]);
		}
		else
		{
			if ( delay > MAXDELAY )
			{
				//print log
				printf("overtime! used %dms, got %d UDP\t", delay, UDPSock[i].gettimes);
				int j = 13;
				while (UDPSock[i].SendBuffer[j] != 0 )
				{
					if (UDPSock[i].SendBuffer[j]<0x20) printf("."); else printf("%c", UDPSock[i].SendBuffer[j]);
					j++;
				}    
				printf("\n");

				CloseSocket(&UDPSock[i]);
			}
		}                 

	}    
}

void ReadResponse()
{
	//read response from dns server
	char                Buffer[1024];
	int                 bytecount;
	int                 fromlen;

	for (int i=0; i<MAXSOCKET; i++)
	{    
		if (UDPSock[i].InUsed && UDPSock[i].UDPSocket!=INVALID_SOCKET && (FD_ISSET(UDPSock[i].UDPSocket, &fds_read)!=0))
		{
			fromlen = sizeof(SOCKADDR_STORAGE);  
			bytecount = recvfrom(UDPSock[i].UDPSocket, Buffer, 1024, 0, (SOCKADDR *)&UDPSock[i].UDPTo, &fromlen);
			UDPSock[i].end = GetTickCount();
			UDPSock[i].gettimes ++;

			if(bytecount>0)
			{
				UDPSock[i].SendBufferLen = bytecount;

				memcpy(UDPSock[i].SendBuffer,Buffer,bytecount);
			}
			else
			{                    
				CloseSocket(&UDPSock[i]);
			}   
		}
	}
}


void ReadRequest()
{
	if (FD_ISSET(localDnsSocket, &fds_read)==0) return;

	char                Buffer[1024];
	SOCKADDR_STORAGE    from;
	int                 fromlen, bytecount, sockNo;

	fromlen = sizeof(SOCKADDR_STORAGE);
	bytecount = recvfrom(localDnsSocket, Buffer, 1024, 0, (SOCKADDR *)&from, &fromlen);

	if (bytecount <= 0) {
		printf("local receve 0 byte!");
		return;
	}
	sockNo = 10000;
	for (int i=0; i<MAXSOCKET; i++)
	{
		if (!UDPSock[i].InUsed)
		{
			UDPSock[i].InUsed=true;
			sockNo = i;
			break;
		}    
	}    
	if (sockNo == 10000){
		printf("Resource poor!");
		return;
	}

	UDPSock[sockNo].UDPFrom = from;
	
	in6_addr    i6addr;
	int ret = CheckIP6( Buffer, &i6addr );

	if (ret==1)
	{
		struct  DNS_HEADER *dns = NULL;
		struct  RES_RECORD  *rr = NULL;
		struct  CNAME_RECORD  *rc = NULL;

		dns = (struct DNS_HEADER *)&Buffer;
		dns->qr = 1;
		dns->ra = 1;
		dns->ans_count_h  = 0;
		dns->ans_count_l  = 1;

		/*rr = (struct RES_RECORD *)&Buffer[bytecount];
		rr->name[0]=(char)0xc0;
		rr->name[1]=(char)0x0c;
		rr->type_h = 0;
		rr->type_l = 28;
		rr->_class_h = 0;
		rr->_class_l = 1;
		rr->ttl_hh = 0x3f;
		rr->ttl_h  = 0xff;
		rr->ttl_l  = 0xff;
		rr->ttl_ll = 0xff;
		rr->data_len_h = 0;
		rr->data_len_l = 16;
		rr->i6addr = i6addr;*/

		rc = (struct CNAME_RECORD *)&Buffer[bytecount];
		rc->name[0]=(char)0xc0;
		rc->name[1]=(char)0x0c;
		rc->type_h = 0;
		rc->type_l = 5;
		rc->_class_h = 0;
		rc->_class_l = 1;
		rc->ttl_hh = 0x3f;
		rc->ttl_h  = 0xff;
		rc->ttl_l  = 0xff;
		rc->ttl_ll = 0xff;
		rc->data_len_h = 0;
		rc->data_len_l = 16;
		rc->cname[0]=3;
		rc->cname[1]='a';
		rc->cname[2]='b';
		rc->cname[3]='c';
		rc->cname[4]=6;
		rc->cname[5]='g';
		rc->cname[6]='o';
		rc->cname[7]='o';
		rc->cname[8]='g';
		rc->cname[9]='l';
		rc->cname[10]='e';
		rc->cname[11]=3;
		rc->cname[12]='c';
		rc->cname[13]='o';
		rc->cname[14]='m';
		rc->cname[15]=0;

		bytecount += 28;

		memcpy(UDPSock[sockNo].SendBuffer,Buffer,bytecount);

		UDPSock[sockNo].SendBufferLen = bytecount;
		UDPSock[sockNo].start = GetTickCount();
		UDPSock[sockNo].end = UDPSock[sockNo].start;
		UDPSock[sockNo].gettimes = 100;
	}    
	else
	{
		memcpy(UDPSock[sockNo].SendBuffer,Buffer,bytecount);
		UDPSock[sockNo].SendBufferLen = bytecount;
		UDPSock[sockNo].UDPSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

		if (UDPSock[sockNo].UDPSocket == INVALID_SOCKET)
		{
			CloseSocket(&UDPSock[sockNo]);
			printf("Server: Error at udp socket(), error code: %ld.\n", WSAGetLastError());
			return;
		}

		if (sendto(UDPSock[sockNo].UDPSocket, UDPSock[sockNo].SendBuffer, UDPSock[sockNo].SendBufferLen, 0, (SOCKADDR *)&UDPSock[sockNo].UDPTo, sizeof(UDPSock[sockNo].UDPTo))==-1)
		{
			printf( "Failed to sendto(), error code: %ld.\n", WSAGetLastError());
			CloseSocket(&UDPSock[sockNo]);
		}
		else
		{
			UDPSock[sockNo].start = GetTickCount();
			UDPSock[sockNo].end = UDPSock[sockNo].start;
		}
	}    
}


void CreateIP6Table ()
{
	FILE    *IP6File;
	char    filename[255], line[255], addr[100], hostname[100];
	int     ret;

	ret = GetCurrentDirectory(255, filename);
	strcat_s ( filename, 255, "\\");
	strcat_s ( filename, 255, IP6HOSTFILE);

	ret = fopen_s( &IP6File, filename, "r");
	if (ret != 0)
	{
		printf("open ip6 hosts file failure!\n");
		return;
	}

	while(fgets(line,sizeof(line),IP6File))
	{
		int index, num, len ;
		num = 0; index = 0;
		len = strlen(line);

		if ( (len>0) &&  (line[0]!='#') && (isspace(line[0])==0) )
		{
			while ( (index<len)&&(isspace(line[index])== 0) ) addr[num++] = line[index++];
			addr[num]='\0';
			while ( (index<len)&&(isspace(line[index])!= 0) ) index++;
			num = 0;        
			while ( (index<len)&&(isspace(line[index])== 0) ) hostname[num++] = line[index++];
			hostname[num]='\0';
			ip6Table.insert( pair <string, string> (hostname, addr));
		}            
	}   
	fclose(IP6File);         
}



int CheckIP6(char* Buffer, in6_addr* i6addr)
{
	int     ret, i, dot;
	char    hostname[100], addr[100];

	i=13; dot=0;
	while ( Buffer[i] != 0 )
	{
		if (Buffer[i]<0x20) 
		{
			hostname[i-13]='.';
			dot++; 
		}
		else hostname[i-13] = Buffer[i];
		i++;
	}
	hostname[i-13]='\0';
	addr[0]='\0';

	//query type =A
	if(Buffer[i+2]!=1 && Buffer[i+2]!=28 && Buffer[i+2]!=255)
	{
		printf("Query Type= %d",Buffer[i+2]);
		return 0;
	}

	//dns suffix
	char * ns=strstr(hostname,DNSSUFFIX);
	if(ns!=NULL)
	{
		ns[0]='\0';
	}

	map<string,string>::iterator p;

	p = ip6Table.find(hostname);
	if (p == ip6Table.end())
	{
		if (dot>1)
		{
			int j;
			i=0; j=0;
			hostname[j]='*'; j++;
			while ( hostname[i]!='.') i++;
			while ( hostname[i]!='\0') hostname[j++]=hostname[i++];
			hostname[j] = '\0';
			p = ip6Table.find(hostname);
		}
		else return 0;        
	}    
	if (p != ip6Table.end())
	{
		strcpy_s(addr, 100, (*p).second.c_str());

		sockaddr_in6    saddr;
		int saddrLen = sizeof(saddr);
		ret = WSAStringToAddress(addr, AF_INET6, NULL, (LPSOCKADDR)&saddr, &saddrLen);
		if (ret != 0) 
		{
			printf("wsastringtoaddress failure! Error code: %ld\n", WSAGetLastError());
			return 0;
		}
		else
		{    
			*i6addr = saddr.sin6_addr;
			return 1;
		}    
	}
	else
	{
		return 0;
	}        
}    
